import argparse
import os
from functools import partial

import torch
import torch.distributed as dist
from datasets import load_dataset
from torch.utils.data import DistributedSampler
from transformers import AutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding
import torch.utils.data as Data
from torch.optim import AdamW
#使用启动工具 torch.distributed.launch 在每个主机上执行一次脚本，开始训练;
#所以使用multiprocessing 方便
import torch.multiprocessing as mp

os.environ[ 'MASTER_ADDR'] = 'localhost'

os.environ[ 'MASTER_PORT'] = '12355'


parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=0, type=int,
                    help='node rank for distributed training')
parser.add_argument('--model_name_or_path', default="bert-base-uncased", type=str,
                    help='node rank for distributed training')
parser.add_argument('--batch_size', default=2, type=int, help='lujing')
parser.add_argument('--num_epochs', default=100, type=int, help='lujing')
parser.add_argument('--save_interval', default=100, type=int, help='lujing')
parser.add_argument('--save_dir', default="./save_model/", type=str, help='lujing')
args = parser.parse_args()
print(args.local_rank)

# 其中 proc 对应 local_rank（当前进程 index），进程数 nprocs 对应 4
def main_worker(proc, nprocs, args):

    # group：进程组，通常一个job只有一个组，即一个world，使用多机时，一个group产生了多个world。
    # world_size：一个job的全局进程数量
    # rank：进程的序号，一般设置rank = 0 的主机为master节点。
    # local_rank：进程内部的GPU序号。
    # 比如，有两台8卡机器，这时具有一个group，2
    # 个world，每个world_size为8，第一个主机rank = 0，显卡编号依次为0, ..., 7，第二个主机rank = 1，显卡编号依次为0, ..., 7。

    """

    1.init_process_group 初始化进程组，同时初始化 distributed 包。
    2.创建分布式模型model = DDP(model)
    3.创建分布式数据采样的datasampler
    4.利用torch.distributed.launch控制进程训练
    5.destory_process_group销毁进程组
    """
    dist.init_process_group(backend='nccl', init_method='tcp://127.0.0.1:23456', world_size=nprocs, rank=proc)

    # 需要对数据集进行划分
    raw_datasets = load_dataset("glue", "mrpc")

    # 首先cpu加载
    model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path,
                                                               device_map='cuda',
                                                               num_labels=raw_datasets.num_columns["train"])
    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)

    tokenized_datasets = raw_datasets.map(lambda data: tokenizer(data["sentence1"], data["sentence2"], padding=True),
                                          batched=True)

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer, padding="max_length", max_length=512, )


    def collate_fn(data, data_collator):
        encodedata = [{'input_ids': indata['input_ids'], 'token_type_ids': indata['token_type_ids'],
                       'attention_mask': indata['attention_mask'], 'label': indata['label']} for indata in data]

        # labels=[indata['label'] for indata in data]
        return data_collator(encodedata)


    collate_fn_partial = partial(collate_fn, data_collator=data_collator)

    # 用 DistributedSampler 对数据集进行划分。如此前我们介绍的那样，它能帮助我们将每个 batch 划分成几个 partition，
    # 在当前进程中只需要获取和 rank 对应的那个 partition 进行训练：
    # 需要更改
    train_dataloader = Data.DataLoader(
        tokenized_datasets["train"], shuffle=False, collate_fn=collate_fn_partial, batch_size=args.batch_size,
        sampler=DistributedSampler(tokenized_datasets["train"])
    )

    #屏幕数据无需sample 有sample  shuffle无用
    dev_dataloader = Data.DataLoader(
        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn_partial, batch_size=args.batch_size,
        sampler=DistributedSampler(tokenized_datasets["validation"])
    )

    torch.cuda.set_device(args.local_rank)

    model.cuda()

    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

    optimizer = AdamW(model.parameters())


    for epoch in range(100):
        print("------------------------the %d epoch-----------------------------------" % epoch)
        for step, batchdata in enumerate(train_dataloader):
            print("----------------------")
            print(model.device)
            print("----------------------")
            # **batchdata ** 用于二维数组  *用于一维 将dict进行拆解，同时会将input_ids等进行参数传递，不用担心匹配出错

            # SequenceClassifierOutput(
            #     loss=loss,
            #     logits=logits,
            #     hidden_states=outputs.hidden_states,
            #     attentions=outputs.attentions,
            # )
            output = model(**batchdata.to(model.device))
            loss = output.loss
            print("----------输出结果------------")
            optimizer.zero_grad()  # 每次计算万的时候需要把上次计算的梯度设置为0
            loss.backward()
            optimizer.step()


mp.spawn(main_worker, nprocs=1, args=(1,args))
